#!coding:utf-8
import numpy as np
from transforms3d.euler import euler2mat
from .config import BehaviorState


class GaitController:
    def __init__(self, config):
        self.config = config

    def phase_index(self, ticks):
        """
        计算在给定时钟周期下所在的相位: 
            同时落地，抬2、3号腿，同时落地，抬1、4号腿
        Parameters
        ----------
        ticks : int
            自程序启动以来的时钟计数

        Returns
        -------
        Int
            当前时钟周期数对应的步态相位
        """
        # 取余得到当前周期的时钟计数
        phase_time = ticks % self.config.phase_length
        phase_sum = 0
        # 累加判断当前时钟数处于哪个周期
        for i in range(self.config.num_phases):
            phase_sum += self.config.phase_ticks[i]
            if phase_time < phase_sum:
                return i
        assert False

    def subphase_ticks(self, ticks):
        """
        计算在当前相位下的时钟计数
        Parameters
        ----------
        ticks : Int
            自程序启动以来的时钟计数

        Returns
        -------
        Int
            当前步态相位的时钟计数
        """
        # 当前周期的时钟计数
        phase_time = ticks % self.config.phase_length
        phase_sum = 0
        subphase_ticks = 0
        for i in range(self.config.num_phases):
            phase_sum += self.config.phase_ticks[i]
            # 找到当前所处的相位
            if phase_time < phase_sum:
                # 计算当前所处的相位的时钟计数
                subphase_ticks = phase_time - \
                    phase_sum + self.config.phase_ticks[i]
                return subphase_ticks
        assert False

    def contacts(self, ticks):
        """
        计算在给定时钟计数下每条腿的状态
        Parameters
        ----------
        ticks : Int
            自程序启动以来的时钟计数

        Returns
        -------
        numpy array (4,)
            numpy数组，0代表抬起，1代表支撑
        """
        # 按列取出contact_phases数组
        return self.config.contact_phases[:, self.phase_index(ticks)]


class StanceController:
    def __init__(self, config):
        self.config = config

    def position_delta(self, leg_index, state, command):
        """
        计算支撑的腿下一时刻的坐标变化（三角形的底）
        Parameters
        ----------
        leg_index : int
            腿的序号

        Returns
        -------
        (Numpy array (3), Numpy array (3, 3))
            (一个时钟周期的坐标变化, 旋转矩阵变化)
        """

        v_xy = np.array(
            [
                # 触地时反向移动
                -command.horizontal_velocity[0],
                -command.horizontal_velocity[1],
                # 暂不考虑身体高度变化（z轴）
                0,
            ]
        )
        delta_p = v_xy * self.config.dt
        return delta_p

    def next_foot_location(self, leg_index, state, command):
        foot_location = state.foot_locations[:, leg_index]
        delta_p = self.position_delta(leg_index, state, command)
        # 转向控制
        delta_R = euler2mat(0, 0, -command.yaw * self.config.dt)
        # 当前坐标加上变化位移就是目标位置
        incremented_location = delta_R.dot(foot_location) + delta_p
        return incremented_location


class SwingController:
    # 计算抬起的腿下一时刻的坐标（三角形的腰）
    def __init__(self, config):
        self.config = config

    # 计算落地点的坐标（只含xy）
    def raibert_touchdown_location(self, leg_index, command):
        delta_p_2d = (
            self.config.alpha
            * self.config.stance_ticks
            * self.config.dt
            * command.horizontal_velocity
        )
        delta_p = np.array([delta_p_2d[0], delta_p_2d[1], 0])
        # 转向控制
        theta = (
            self.config.beta
            * self.config.stance_ticks
            * self.config.dt
            * command.yaw
        )
        R = euler2mat(0, 0, theta)
        # 交换旋转和平移的顺序影响不大
        return R.dot(self.config.default_stance[:, leg_index] + delta_p)

    # 计算抬腿高度的变化
    def swing_height(self, swing_phase, triangular=True):
        if triangular:
            # 三角步态，前半段上升，后半段下降
            if swing_phase < 0.5:
                swing_height_ = swing_phase / 0.5 * self.config.z_clearance
            else:
                swing_height_ = self.config.z_clearance * \
                    (1 - (swing_phase - 0.5) / 0.5)
        return swing_height_

    def next_foot_location(
        self,
        swing_prop,
        leg_index,
        state,
        command,
    ):
        assert swing_prop >= 0 and swing_prop <= 1
        # 当前时刻的坐标
        foot_location = state.foot_locations[:, leg_index]
        # 抬腿周期的高度
        swing_height_ = self.swing_height(swing_prop)
        # 计算落地点的坐标（只含xy）
        touchdown_location = self.raibert_touchdown_location(
            leg_index, command)
        # 计算抬腿的剩余时间
        time_left = self.config.dt * \
            self.config.swing_ticks * (1.0 - swing_prop)
        # 计算当前速度
        v = (touchdown_location - foot_location) / \
            time_left * np.array([1, 1, 0])
        # 计算水平位移
        delta_foot_location = v * self.config.dt
        # 计算z轴
        z_vector = np.array([0, 0, swing_height_ + command.height])
        # 最后加起来，考虑到可能在任意时刻改变速度和方向，这样计算可以及时响应
        return foot_location * np.array([1, 1, 0]) + z_vector + delta_foot_location


class Controller:
    """
    主控制器，一步一步计算每只脚的位置
    """

    def __init__(
        self,
        config,
    ):
        self.config = config
        self.contact_modes = np.zeros(4)
        self.gait_controller = GaitController(self.config)
        self.swing_controller = SwingController(self.config)
        self.stance_controller = StanceController(self.config)

        # 状态机转换
        self.trot_transition_mapping = {
            BehaviorState.DEACTIVATED: BehaviorState.REST, BehaviorState.REST: BehaviorState.TROT, BehaviorState.TROT: BehaviorState.REST}
        self.activate_transition_mapping = {
            BehaviorState.DEACTIVATED: BehaviorState.REST, BehaviorState.REST: BehaviorState.DEACTIVATED, BehaviorState.TROT: BehaviorState.REST}

    def step_gait(self, state, command):
        """
            计算下一时刻的四足位置
        Returns
        -------
        Numpy array (3, 4)
            表示四足位置的矩阵
        """
        # 获取四足与地面的接触状态
        contact_modes = self.gait_controller.contacts(state.ticks)
        new_foot_locations = np.zeros((3, 4))
        # 分别计算四条腿的位置
        for leg_index in range(4):
            contact_mode = contact_modes[leg_index]
            # 触地阶段，z轴不变，xy往速度的反方向运动，直线，三角形的底
            if contact_mode == 1:
                new_location = self.stance_controller.next_foot_location(
                    leg_index, state, command)
            # 抬腿阶段，z先上升再下降，xy往速度的方向运动，尖，三角形的腰
            else:
                # 抬腿的进度，0表示开始（将要抬起），1表示结束（放下落地）
                swing_proportion = (
                    self.gait_controller.subphase_ticks(
                        state.ticks)*1.0 / self.config.swing_ticks
                )
                new_location = self.swing_controller.next_foot_location(
                    swing_proportion,
                    leg_index,
                    state,
                    command
                )
            new_foot_locations[:, leg_index] = new_location
        return new_foot_locations

    def run(self, state, command):
        """Steps the controller forward one timestep
        使控制器前进一个时钟周期
        Parameters
        ----------
        state : State
            机器人当前状态
        command : Command
            机器人控制指令            
        """

        if command.activate_event:
            state.behavior_state = self.activate_transition_mapping[state.behavior_state]
        elif command.trot_event:
            state.behavior_state = self.trot_transition_mapping[state.behavior_state]

        if state.behavior_state == BehaviorState.REST:
            state.foot_locations = (
                self.config.default_stance
                + np.array([0, 0, command.height])[:, np.newaxis]
            )
            # 乘以旋转矩阵
            state.rotated_foot_locations = (
                euler2mat(
                    command.roll,
                    command.pitch,
                    command.yaw/3.0,
                ).dot(state.foot_locations)
            )
        elif state.behavior_state == BehaviorState.TROT:
            state.foot_locations = self.step_gait(
                state,
                command,
            )

            # state.foot_locations会被用于StanceController的位移计算
            # 所以这里另外定义rotated_foot_locations计算旋转后的坐标防止多次计算
            state.rotated_foot_locations = (
                euler2mat(
                    command.roll,
                    command.pitch,
                    0,
                ).dot(state.foot_locations)
            )
        # 时钟计数+1
        state.ticks += 1
        # 更新机器人状态
        if abs(command.roll) <= self.config.max_roll:
            state.roll = command.roll
        if abs(command.pitch) <= self.config.max_pitch:
            state.pitch = command.pitch
        if abs(command.yaw) <= self.config.max_yaw:
            state.yaw = command.yaw
        state.height = command.height
